from typing import Union

import numpy as np
import tensorflow as tf
from scipy.sparse import issparse, spmatrix, coo_matrix, csr_matrix
from tensorflow import sparse

spt = sparse.SparseTensor


def to_dense_tensor(mat: Union[spmatrix, np.ndarray, dict], dtype=None) -> tf.Tensor:
    mat = to_sparse_tensor(mat, dtype)
    mat = tf.sparse.reorder(mat)
    return sparse.to_dense(mat)


def to_sparse_tensor(matrix: Union[spmatrix, np.ndarray, dict], dtype=None) -> spt:
    if isinstance(matrix, dict):
        matrix = sparse_from_json(matrix)
        matrix: csr_matrix
    if issparse(matrix):
        matrix = matrix.tocoo()
        matrix: coo_matrix
        row = matrix.row
        col = matrix.col
        values = matrix.data
    else:
        matrix: np.ndarray
        row, col = matrix.nonzero()
        values = matrix[row, col]
    indices = np.array([row, col]).T
    if dtype is not None and (isinstance(dtype, str) or isinstance(dtype, np.dtype)):
        values = values.astype(dtype)
    if dtype is not None and isinstance(dtype, tf.dtypes.DType):
        values = tf.constant(values, dtype=dtype)
    return spt(indices=indices, values=values, dense_shape=matrix.shape)


def from_sparse_tensor(tensor: spt, dtype=None) -> csr_matrix:
    data = tensor.values.numpy()
    indices = tensor.indices.numpy()
    row, col = indices.T
    if dtype is not None:
        data = data.astype(dtype)
    return csr_matrix((data, (row, col)), shape=tensor.dense_shape)


def sparse_to_json(mat: Union[spt, spmatrix]) -> dict:
    if isinstance(mat, spt):
        mat = from_sparse_tensor(mat)
    mat: csr_matrix = mat.tocsr()
    return {
        'indptr': mat.indptr,
        'indices': mat.indices,
        'data': mat.data,
        'shape': mat.shape,
        'dtype': str(mat.dtype)
    }


def sparse_from_json(mat_json: dict) -> csr_matrix:
    data = mat_json['data']
    indices = mat_json['indices']
    indptr = mat_json['indptr']
    m, n = mat_json['shape']
    dtype = mat_json['dtype']
    return csr_matrix((data, indices, indptr), shape=(m, n), dtype=dtype)


def save_dense_in_sparse_json(mat: tf.Tensor) -> dict:
    return sparse_to_json(sparse.from_dense(mat))
